import json
import os
import xml.etree.ElementTree as ET
from xml.dom import minidom

import random
from collections import defaultdict
def json_to_xml(json_path, output_dir):
    # 创建输出目录
    os.makedirs(output_dir, exist_ok=True)

    # 读取JSON文件
    with open(json_path, 'r') as f:
        data = json.load(f)

    # 构建类别ID到名称的映射
    categories = {cat['id']: cat['name'] for cat in data['categories']}

    # 按图像ID分组标注
    image_anns = {}
    for ann in data['annotations']:
        image_id = ann['image_id']
        if image_id not in image_anns:
            image_anns[image_id] = []
        image_anns[image_id].append(ann)

    # 构建图像ID到图像信息的映射
    images_info = {img['id']: img for img in data['images']}

    # 为每张图像生成XML
    for img_id, anns in image_anns.items():
        img_info = images_info[img_id]
        filename = img_info['file_name']
        width = img_info['width']
        height = img_info['height']

        # 创建XML根节点
        root = ET.Element('annotation')

        # 添加基本信息
        ET.SubElement(root, 'folder').text = 'VOC'
        ET.SubElement(root, 'filename').text = filename

        # 添加图像尺寸
        size = ET.SubElement(root, 'size')
        ET.SubElement(size, 'width').text = str(width)
        ET.SubElement(size, 'height').text = str(height)
        ET.SubElement(size, 'depth').text = '3'  # 假设RGB图像

        ET.SubElement(root, 'segmented').text = '0'

        # 添加每个检测对象
        for ann in anns:
            obj = ET.SubElement(root, 'object')
            ET.SubElement(obj, 'name').text = categories[ann['category_id']]
            ET.SubElement(obj, 'pose').text = 'Unspecified'
            ET.SubElement(obj, 'truncated').text = '0'
            ET.SubElement(obj, 'difficult').text = '0'

            # 转换边界框格式: [x_min, y_min, width, height] -> [xmin, ymin, xmax, ymax]
            bbox = ann['bbox']
            xmin = bbox[0]
            ymin = bbox[1]
            xmax = xmin + bbox[2]
            ymax = ymin + bbox[3]

            # 确保坐标不超出图像范围
            xmin = max(0, min(xmin, width))
            ymin = max(0, min(ymin, height))
            xmax = max(0, min(xmax, width))
            ymax = max(0, min(ymax, height))

            # 添加边界框
            bndbox = ET.SubElement(obj, 'bndbox')
            ET.SubElement(bndbox, 'xmin').text = str(int(xmin))
            ET.SubElement(bndbox, 'ymin').text = str(int(ymin))
            ET.SubElement(bndbox, 'xmax').text = str(int(xmax))
            ET.SubElement(bndbox, 'ymax').text = str(int(ymax))

        # 生成XML字符串并美化输出
        xml_str = ET.tostring(root, encoding='utf-8')
        dom = minidom.parseString(xml_str)
        pretty_xml = dom.toprettyxml(indent='  ')

        # 写入文件 (使用图像ID作为文件名)
        xml_filename = os.path.splitext(filename)[0] + '.xml'
        xml_path = os.path.join(output_dir, xml_filename)
        with open(xml_path, 'w') as xml_file:
            xml_file.write(pretty_xml)



def write_filenames_without_ext(folder_path, output_file):
    """
    将文件夹内所有文件的文件名（不带后缀）写入到文本文件中

    参数:
        folder_path: 要扫描的文件夹路径
        output_file: 输出的文本文件路径
    """
    # 获取文件夹内所有文件名（不包括子目录）
    filenames = [f for f in os.listdir(folder_path)
                 if os.path.isfile(os.path.join(folder_path, f))]

    # 去掉文件后缀
    names_without_ext = [os.path.splitext(f)[0] for f in filenames]

    # 写入到文本文件
    with open(output_file, 'w') as f:
        for name in names_without_ext:
            f.write(name + '\n')

    print(f"成功写入 {len(names_without_ext)} 个文件名到 {output_file}")


def split_images_by_category(input_txt, train_txt, val_txt, train_ratio=0.7):
    # 读取输入文件中的所有图片名称
    with open(input_txt, 'r') as f:
        image_names = [line.strip() for line in f.readlines()]

    # 按类别分组图片
    category_images = defaultdict(list)
    for name in image_names:
        # 提取类别（第一个下划线前的部分）
        if '_' in name:
            category = name.split('_', 1)[0]
            category_images[category].append(name)
        else:
            # 没有下划线的图片处理
            print(f"警告: 图片名称 '{name}' 不包含下划线，无法提取类别")
            category_images[name].append(name)  # 使用完整名称作为类别

    # 计算总图片数和类别数
    total_images = len(image_names)
    total_categories = len(category_images)
    print(f"总类别数: {total_categories}, 总图片数: {total_images}")

    # 将类别随机排序
    categories = list(category_images.keys())
    random.shuffle(categories)

    # 初始化训练集和验证集
    train_set = []
    val_set = []
    train_count = 0
    val_count = 0
    train_categories = []
    val_categories = []

    # 计算目标训练集大小
    target_train_size = int(total_images * train_ratio)

    # 分配类别到训练集或验证集
    for category in categories:
        images = category_images[category]
        num_images = len(images)

        # 如果当前训练集图片数 + 该类图片数 <= 目标大小，则加入训练集
        if train_count + num_images <= target_train_size:
            train_set.extend(images)
            train_count += num_images
            train_categories.append(category)
        else:
            # 如果加入训练集会超出目标大小，则加入验证集
            val_set.extend(images)
            val_count += num_images
            val_categories.append(category)

    # 计算最终比例
    actual_ratio = train_count / total_images
    print(f"\n训练集: {len(train_set)} 张图片 ({len(train_categories)} 个类别)")
    print(f"验证集: {len(val_set)} 张图片 ({len(val_categories)} 个类别)")
    print(f"实际分割比例: 训练集 {actual_ratio:.2%}, 验证集 {1 - actual_ratio:.2%}")

    # 写入训练集文件
    with open(train_txt, 'w') as f:
        for name in train_set:
            f.write(name + '\n')

    # 写入验证集文件
    with open(val_txt, 'w') as f:
        for name in val_set:
            f.write(name + '\n')

    # 返回划分结果
    return {
        "train_count": len(train_set),
        "val_count": len(val_set),
        "train_categories": train_categories,
        "val_categories": val_categories
    }



if __name__ == "__main__":
    # # 示例用法
    # json_path = r"D:\Code\0-data\1-齿轮检测数据集\青山数据集\qingshan_dianquketi_v1.json"  # 输入JSON文件路径
    # output_dir = r"D:\Code\0-data\1-齿轮检测数据集\青山数据集\xml"  # 输出XML目录
    #
    # json_to_xml(json_path, output_dir)
    # print(f"转换完成！XML文件已保存至: {output_dir}")
    # 配置参数
    # input_folder = r"D:\Code\0-data\1-齿轮检测数据集\青山数据集\xml"  # 替换为你的文件夹路径
    # output_txt = r"D:\Code\0-data\1-齿轮检测数据集\青山数据集\all.txt"  # 输出的文本文件名
    #
    # write_filenames_without_ext(input_folder, output_txt)
    input_txt = r"D:\Code\0-data\1-齿轮检测数据集\青山数据集\VOCdevkit\VOC2012\ImageSets\Main\extra_images.txt"  # 输入文件：包含所有图片名称（不带后缀）
    train_txt = r"D:\Code\0-data\1-齿轮检测数据集\青山数据集\VOCdevkit\VOC2012\ImageSets\Main\ex_train.txt"  # 输出文件：训练集图片名称
    val_txt = r"D:\Code\0-data\1-齿轮检测数据集\青山数据集\VOCdevkit\VOC2012\ImageSets\Main\ex_val.txt"  # 输出文件：验证集图片名称

    # 设置随机种子确保可重复性
    random.seed(0)

    # 执行划分
    result = split_images_by_category(input_txt, train_txt, val_txt)

    # 打印类别划分详情
    print("\n=== 训练集类别 ===")
    print(", ".join(sorted(result["train_categories"])))

    print("\n=== 验证集类别 ===")
    print(", ".join(sorted(result["val_categories"])))